CPU


In [ ]:
import numpy as np
import tensorflow as tf
import time

with tf.device('/cpu:0'):
  X = tf.constant(np.array(np.random.randn(10000,10000), dtype = np.float32), dtype = tf.float32)
  Y = tf.matmul(X, tf.transpose(X))

  init = tf.initialize_all_variables()
  sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
  sess.run(init)

  a = time.time()
  sess.run(Y)
  print "Runtime: %s" % (time.time() - a)

Single GPU


In [ ]:
import tensorflow as tf
import time
import numpy as np

# Creates a graph.
c = []
for d in ['/gpu:0', '/gpu:1', '/gpu:2', '/gpu:3']:
#for d in ['/gpu:0']:
  with tf.device(d):
    a = tf.constant(np.array(np.random.randn(500,500), dtype = np.float32), dtype = tf.float32)
    b = tf.constant(np.array(np.random.randn(500,500), dtype = np.float32), dtype = tf.float32)
    c = tf.constant(np.array(np.random.randn(500,500), dtype = np.float32), dtype = tf.float32)
    d = tf.constant(np.array(np.random.randn(500,500), dtype = np.float32), dtype = tf.float32)
    c.append(tf.matmul(a, b))
    c.append(tf.matmul(c, d))

with tf.device('/cpu:0'):
  sum = tf.add_n(c)

# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# Runs the op.

a = time.time()
print sess.run(sum)
print "Runtime: %s" % (time.time() - a)

sess.close()

Multiple GPUs


In [ ]:
import tensorflow as tf
import time
import numpy as np

# Creates a graph.
c = []
for d in ['/gpu:0', '/gpu:1', '/gpu:2', '/gpu:3']:
#for d in ['/gpu:0']:
  with tf.device(d):
    a = tf.constant(np.array(np.random.randn(500,500), dtype = np.float32), dtype = tf.float32)
    b = tf.constant(np.array(np.random.randn(500,500), dtype = np.float32), dtype = tf.float32)
    c = tf.constant(np.array(np.random.randn(500,500), dtype = np.float32), dtype = tf.float32)
    d = tf.constant(np.array(np.random.randn(500,500), dtype = np.float32), dtype = tf.float32)
    c.append(tf.matmul(a, b))
    c.append(tf.matmul(c, d))

with tf.device('/cpu:0'):
  sum = tf.add_n(c)

# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# Runs the op.

a = time.time()
print sess.run(sum)
print "Runtime: %s" % (time.time() - a)

sess.close()

In [ ]: